/*
* Copyright (C) Microsoft. All rights reserved.  
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not 
* use this file except in compliance with the License.  You may obtain a copy 
* of the License at http://www.apache.org/licenses/LICENSE-2.0  
* 
* THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED 
* WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, 
* MERCHANTABLITY OR NON-INFRINGEMENT. 
*
* See the Apache License, Version 2.0 for specific language governing 
* permissions and limitations under the License.
*
*/
#pragma once

#include <baja/common.hpp>
#include <map>
#include <vector>
#include <limits.h>
#include <functional>

namespace baja
{

template <class T>
class heap
{
public:
    typedef size_t handle;

    heap(
        const size_t capacity,
        const std::function<int(const T&, const T&)> comparison = 
            std::bind(&heap<T>::defaultComparison, std::placeholders::_1, std::placeholders::_2)
        ) : 
        _capacity(capacity), _comparison(comparison), _size(0), _lastHandle(0)
    {
        _data.resize(capacity);
    }

    handle push(const T& data)
    {
        if (full())
            throw std::logic_error("heap is full");

        handle h = newHandle();
        _data[_size].data = data;
        _data[_size].h = h;

        size_t i = bubble(_size);
        _size++;

        _handles[h] = i;
        return h;
    }

    void touch(const handle h, std::function<void(T&)> t)
    {
        size_t n = _handles[h];

        t(_data[n].data);
        n = sink(bubble(n));

        _handles[h] = n;
    }

    T pop()
    {
        if (size() == 0) throw std::logic_error("heap is empty");

        T t = _data[0].data;
        handle handleToDelete = _data[0].h;

        _size--;
        if (size() > 0)
        {
            swap(0, _size);
            sink(0);
        }
        // erase() needs to be called after swap().
        // If not, swap() will re-create an entry for _handles
        _handles.erase(handleToDelete);
        return t;
    }

    void remove(const handle h)
    {
        if (size() == 0) throw std::logic_error("heap is empty");

        size_t n = _handles[h];
        
        _size--;
        if (size() > 0)
        {
            swap(n, _size);
            sink(bubble(n));
        }
        // erase() needs to be called after swap().
        // If not, swap() will re-create an entry for _handles
        _handles.erase(h);
    }

    T front() const { return _data[0].data; }

    size_t size() const { return _size; }
    size_t capacity() const { return _capacity; }
    bool full() const { return _size >= _capacity; }
    bool empty() const { return _size == 0; }
    handle newHandle() { return ++_lastHandle; }

    inline const size_t invalid() const { return SIZE_MAX; }

private:
    static int defaultComparison(const T& left, const T& right)
    {
        if (left < right) return -1;
        else if (left > right) return 1;
        else return 0;
    }

    void swap(const size_t i, const size_t j)
    {
        heap_data temp = _data[i];
        _data[i] = _data[j];
        _data[j] = temp;
        _handles[_data[i].h] = i;
        _handles[_data[j].h] = j;
    }

    int compare(const size_t i, const size_t j)
    {
        return _comparison(_data[i].data, _data[j].data);
    }

    size_t left(const size_t i) 
    { 
        size_t j = 2*i + 1; 
        return (j < _size) ? j : invalid();
    }

    size_t right(const size_t i) 
    { 
        size_t j = 2*i + 2;
        return (j < _size) ? j : invalid(); 
    }

    size_t parent(const size_t i)
    {
        if (i == 0) return invalid();
        return static_cast<size_t>(floor((i - 1) / 2));
    }

    size_t bubble(const size_t i)
    {
        size_t n = i;
        size_t j = parent(n);
        while (j != invalid())
        {
            if (compare(n, j) > 0)
            {
                // parent is smaller
                break;
            }
            swap(n, j);
            n = j;
            j = parent(n);
        }

        return n;
    }

    size_t sink(const size_t i)
    {
        size_t n = i;
        size_t j = left(n);

        while (j != invalid())
        {
            if (compare(n, j) > 0)
            {
                // left child is smaller
                size_t k = right(n);
                if (k != invalid() && compare(j, k) > 0)
                {
                    // right child is even smaller
                    j = k;
                }
            }
            else
            {
                break;
            }

            swap(n, j);
            n = j;
            j = left(n);
        }

        return n;
    }

private:
    struct heap_data
    {
        handle h;
        T data;
    };

    size_t _capacity;
    size_t _size;
    std::function<int(const T&, const T&)> _comparison;
    std::vector<heap_data> _data;
    std::map<handle, size_t> _handles;
    handle _lastHandle;
};

} // namespace baja